#!/usr/bin/env python
# coding: utf-8

# In[1]:


import tushare as ts

# 设置Tushare token（需替换为个人token）
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()


# In[2]:


def get_stock_data(ts_code, start_date, end_date):
    """获取股票日线行情数据"""
    try:
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        # 按交易日期升序排列
        df = df.sort_values('trade_date')
        return df
    except Exception as e:
        print(f"获取{ts_code}日线数据出错: {e}")
        return None

def get_financial_data(ts_code):
    """获取股票财务数据"""
    financial_data = {}
    
    # 获取主要财务指标
    try:
        profit = pro.fina_indicator(ts_code=ts_code)
        financial_data['profit'] = profit
    except Exception as e:
        print(f"获取{ts_code}财务指标出错: {e}")
    
    # 获取资产负债表数据
    try:
        balance = pro.balancesheet(ts_code=ts_code)
        financial_data['balance'] = balance
    except Exception as e:
        print(f"获取{ts_code}资产负债表数据出错: {e}")
    
    # 获取利润表数据
    try:
        income = pro.income(ts_code=ts_code)
        financial_data['income'] = income
    except Exception as e:
        print(f"获取{ts_code}利润表数据出错: {e}")
    
    # 获取现金流量表数据
    try:
        cashflow = pro.cashflow(ts_code=ts_code)
        financial_data['cashflow'] = cashflow
    except Exception as e:
        print(f"获取{ts_code}现金流量表数据出错: {e}")
    
    return financial_data


# In[3]:


def clean_data(df):
    """清洗股票日线数据"""
    if df is None or df.empty:
        return None
    
    # 删除重复记录
    df = df.drop_duplicates()
    
    # 处理缺失值
    df = df.fillna(method='ffill')  # 前向填充
    df = df.dropna()  # 删除仍有缺失值的行
    
    # 处理异常值
    # 剔除涨跌幅超过10%的记录（非ST股票）
    if 'pct_chg' in df.columns and 'is_st' not in df.columns:
        df = df[(df['pct_chg'] <= 10) & (df['pct_chg'] >= -10)]
    
    # 剔除成交量为0的记录
    if 'vol' in df.columns:
        df = df[df['vol'] > 0]
    
    return df

def clean_financial_data(financial_data):
    """清洗财务数据"""
    cleaned_data = {}
    
    for key, df in financial_data.items():
        if df is not None and not df.empty:
            # 删除重复记录
            df = df.drop_duplicates()
            # 处理缺失值
            df = df.fillna(method='ffill')
            df = df.dropna()
            cleaned_data[key] = df
    
    return cleaned_data


# In[4]:


def align_financial_data(daily_df, financial_df):
    """将财务数据与日线数据时间对齐"""
    if daily_df is None or financial_df is None:
        return None
    
    # 确保财务数据按报告期升序排列
    financial_df = financial_df.sort_values('end_date')
    
    # 创建时间对齐后的财务数据DataFrame
    aligned_df = pd.DataFrame(index=daily_df.index)
    
    # 对每个财务指标进行时间对齐
    for col in financial_df.columns:
        if col == 'end_date':
            continue
        
        # 创建一个Series，索引为报告期，值为对应的财务指标
        values = pd.Series(financial_df[col].values, index=pd.to_datetime(financial_df['end_date']))
        
        # 使用ffill方法将财务指标填充到后续交易日
        aligned_values = values.reindex(aligned_df.index, method='ffill')
        
        # 将对齐后的财务指标添加到结果DataFrame
        aligned_df[col] = aligned_values
    
    return aligned_df


# In[5]:


def feature_engineering(df):
    """特征工程：计算技术指标和其他特征"""
    if df is None or df.empty:
        return None
    
    # 复制原始数据
    result = df.copy()
    
    # 计算移动平均线
    result['MA5'] = result['close'].rolling(window=5).mean()
    result['MA10'] = result['close'].rolling(window=10).mean()
    result['MA20'] = result['close'].rolling(window=20).mean()
    result['MA60'] = result['close'].rolling(window=60).mean()
    
    # 计算相对强弱指标(RSI)
    delta = result['close'].diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)
    avg_gain = gain.rolling(window=14).mean()
    avg_loss = loss.rolling(window=14).mean()
    rs = avg_gain / avg_loss
    result['RSI'] = 100 - (100 / (1 + rs))
    
    # 计算布林带
    middle_band = result['close'].rolling(window=20).mean()
    std = result['close'].rolling(window=20).std()
    upper_band = middle_band + 2 * std
    lower_band = middle_band - 2 * std
    result['BB_Middle'] = middle_band
    result['BB_Upper'] = upper_band
    result['BB_Lower'] = lower_band
    result['BB_Width'] = (upper_band - lower_band) / middle_band
    
    # 计算MACD指标
    exp1 = result['close'].ewm(span=12, adjust=False).mean()
    exp2 = result['close'].ewm(span=26, adjust=False).mean()
    result['MACD'] = exp1 - exp2
    result['MACD_Signal'] = result['MACD'].ewm(span=9, adjust=False).mean()
    result['MACD_Hist'] = result['MACD'] - result['MACD_Signal']
    
    # 提取时间特征
    result['trade_date'] = pd.to_datetime(result.index)
    result['day_of_week'] = result['trade_date'].dt.dayofweek
    result['month'] = result['trade_date'].dt.month
    result['quarter'] = result['trade_date'].dt.quarter
    
    # 计算每日涨跌幅
    result['return'] = result['close'].pct_change()
    
    # 删除辅助列
    result = result.drop(columns=['trade_date'])
    
    return result


# In[6]:


import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
from matplotlib.dates import DateFormatter
import matplotlib.ticker as mticker

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
sns.set_style('whitegrid')

class QuantVisualizer:
    """量化交易可视化工具类"""
    
    def __init__(self):
        pass
    
    def plot_stock_price(self, df, ts_code, title="股票价格走势"):
        """绘制股票价格走势图"""
        fig, ax = plt.subplots(figsize=(14, 7))
        
        # 绘制收盘价和均线
        ax.plot(df.index, df['close'], label='收盘价', linewidth=2)
        if 'MA5' in df.columns:
            ax.plot(df.index, df['MA5'], label='MA5', linewidth=1)
        if 'MA20' in df.columns:
            ax.plot(df.index, df['MA20'], label='MA20', linewidth=1)
        if 'MA60' in df.columns:
            ax.plot(df.index, df['MA60'], label='MA60', linewidth=1)
        
        # 添加成交量子图
        ax2 = ax.twinx()
        ax2.bar(df.index, df['vol'], alpha=0.3, color='gray', label='成交量')
        ax2.set_ylabel('成交量')
        
        # 设置标题和标签
        ax.set_title(f"{ts_code} {title}", fontsize=16)
        ax.set_xlabel('日期')
        ax.set_ylabel('价格')
        ax.legend(loc='upper left')
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        ax.xaxis.set_major_formatter(date_format)
        fig.autofmt_xdate()
        
        plt.tight_layout()
        return fig
    
    def plot_technical_indicators(self, df, ts_code):
        """绘制技术指标图"""
        fig, axes = plt.subplots(3, 1, figsize=(14, 12), sharex=True)
        
        # 1. 价格和布林带
        axes[0].plot(df.index, df['close'], label='收盘价')
        axes[0].plot(df.index, df['BB_Upper'], 'g--', label='布林带上轨')
        axes[0].plot(df.index, df['BB_Middle'], 'y--', label='布林带中轨')
        axes[0].plot(df.index, df['BB_Lower'], 'r--', label='布林带下轨')
        axes[0].fill_between(df.index, df['BB_Lower'], df['BB_Upper'], alpha=0.1)
        axes[0].set_title(f"{ts_code} 价格与布林带")
        axes[0].legend()
        
        # 2. RSI指标
        axes[1].plot(df.index, df['RSI'])
        axes[1].axhline(y=70, color='r', linestyle='-', alpha=0.3)
        axes[1].axhline(y=30, color='g', linestyle='-', alpha=0.3)
        axes[1].set_title("RSI指标")
        axes[1].set_ylim(0, 100)
        
        # 3. MACD指标
        axes[2].plot(df.index, df['MACD'], label='MACD')
        axes[2].plot(df.index, df['MACD_Signal'], label='Signal')
        axes[2].bar(df.index, df['MACD_Hist'], label='Histogram')
        axes[2].set_title("MACD指标")
        axes[2].legend()
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        axes[2].xaxis.set_major_formatter(date_format)
        fig.autofmt_xdate()
        
        plt.tight_layout()
        return fig
    
    def plot_financial_indicators(self, df, ts_code):
        """绘制财务指标图"""
        if 'roe' not in df.columns or 'net_profit_margin' not in df.columns:
            print("数据中不包含ROE或净利润率指标")
            return None
        
        fig, axes = plt.subplots(2, 1, figsize=(14, 10), sharex=True)
        
        # 1. ROE和净利润率
        axes[0].plot(df.index, df['roe'], label='ROE(%)', color='blue')
        axes[0].set_title(f"{ts_code} ROE和净利润率")
        axes[0].set_ylabel('ROE(%)')
        axes[0].legend(loc='upper left')
        
        ax2 = axes[0].twinx()
        ax2.plot(df.index, df['net_profit_margin'], label='净利润率(%)', color='orange')
        ax2.set_ylabel('净利润率(%)')
        ax2.legend(loc='upper right')
        
        # 2. 资产负债率和流动比率
        if 'debt_to_assets' in df.columns and 'current_ratio' in df.columns:
            axes[1].plot(df.index, df['debt_to_assets'], label='资产负债率(%)', color='green')
            axes[1].set_title(f"{ts_code} 资产负债率和流动比率")
            axes[1].set_ylabel('资产负债率(%)')
            axes[1].legend(loc='upper left')
            
            ax3 = axes[1].twinx()
            ax3.plot(df.index, df['current_ratio'], label='流动比率', color='red')
            ax3.set_ylabel('流动比率')
            ax3.legend(loc='upper right')
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        axes[1].xaxis.set_major_formatter(date_format)
        fig.autofmt_xdate()
        
        plt.tight_layout()
        return fig
    
    def plot_correlation_matrix(self, df, title="特征相关性矩阵"):
        """绘制特征相关性矩阵"""
        # 选择数值列
        numeric_cols = df.select_dtypes(include=[np.number]).columns
        
        # 计算相关系数矩阵
        corr = df[numeric_cols].corr()
        
        # 创建掩码以只显示下三角
        mask = np.triu(np.ones_like(corr, dtype=bool))
        
        # 设置图形大小
        plt.figure(figsize=(14, 12))
        
        # 绘制热力图
        sns.heatmap(corr, mask=mask, annot=True, fmt=".2f", cmap='coolwarm', 
                    square=True, linewidths=.5, cbar_kws={"shrink": .8})
        
        plt.title(title, fontsize=16)
        plt.tight_layout()
        return plt.gcf()
    
    def plot_strategy_returns(self, results_dict, title="策略收益对比"):
        """绘制不同策略的收益曲线对比图"""
        plt.figure(figsize=(14, 8))
        
        for strategy_name, result_df in results_dict.items():
            plt.plot(result_df.index, result_df['cumulative_return'], 
                     label=strategy_name, linewidth=2)
        
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.title(title, fontsize=16)
        plt.xlabel('日期')
        plt.ylabel('累计收益率')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_strategy_performance(self, result_df, strategy_name):
        """绘制策略绩效综合图"""
        fig, axes = plt.subplots(2, 1, figsize=(14, 12), sharex=True)
        
        # 1. 权益曲线
        axes[0].plot(result_df.index, result_df['total'], label='总资产')
        axes[0].set_title(f"{strategy_name} 策略权益曲线", fontsize=14)
        axes[0].set_ylabel('资产价值(元)')
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)
        
        # 2. 每日收益率
        axes[1].bar(result_df.index, result_df['return'], alpha=0.6)
        axes[1].axhline(y=0, color='black', linestyle='-', alpha=0.3)
        axes[1].set_title(f"{strategy_name} 每日收益率", fontsize=14)
        axes[1].set_ylabel('收益率')
        axes[1].grid(True, alpha=0.3)
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        axes[1].xaxis.set_major_formatter(date_format)
        fig.autofmt_xdate()
        
        # 设置y轴为百分比格式
        axes[1].yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        plt.tight_layout()
        return fig
    
    def plot_drawdown(self, result_df, strategy_name):
        """绘制回撤曲线"""
        # 计算累计收益和最大回撤
        result_df['cumulative_return'] = (1 + result_df['return']).cumprod() - 1
        result_df['peak'] = result_df['cumulative_return'].cummax()
        result_df['drawdown'] = (result_df['cumulative_return'] - result_df['peak']) / (1 + result_df['peak'])
        
        plt.figure(figsize=(14, 7))
        
        # 绘制累计收益曲线
        plt.subplot(2, 1, 1)
        plt.plot(result_df.index, result_df['cumulative_return'], label='累计收益率')
        plt.plot(result_df.index, result_df['peak'], label='历史最高点', linestyle='--', alpha=0.7)
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.title(f"{strategy_name} 累计收益率与最大回撤", fontsize=14)
        plt.ylabel('收益率')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 绘制回撤曲线
        plt.subplot(2, 1, 2)
        plt.fill_between(result_df.index, result_df['drawdown'], 0, color='red', alpha=0.3)
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.5)
        plt.ylabel('回撤')
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        plt.gca().xaxis.set_major_formatter(date_format)
        plt.gcf().autofmt_xdate()
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_trades(self, df, signals, ts_code, strategy_name):
        """绘制交易信号图"""
        plt.figure(figsize=(14, 8))
        
        # 绘制收盘价
        plt.plot(df.index, df['close'], label='收盘价', linewidth=2)
        
        # 标记买入点
        buy_points = signals[signals['signal'] == 1]
        plt.scatter(buy_points.index, buy_points['close'], marker='^', 
                   color='g', s=100, label='买入信号')
        
        # 标记卖出点
        sell_points = signals[signals['signal'] == -1]
        plt.scatter(sell_points.index, sell_points['close'], marker='v', 
                   color='r', s=100, label='卖出信号')
        
        plt.title(f"{ts_code} {strategy_name} 交易信号图", fontsize=16)
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        plt.gca().xaxis.set_major_formatter(date_format)
        plt.gcf().autofmt_xdate()
        
        plt.tight_layout()
        return plt.gcf()


# In[11]:


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.dates import DateFormatter
import matplotlib.ticker as mticker
import tushare as ts

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
sns.set_style('whitegrid')

# 设置Tushare token（需替换为个人token）
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

class QuantVisualizer:
    """量化交易可视化工具类"""
    
    def __init__(self):
        pass
    
    def plot_stock_price(self, df, ts_code, title="股票价格走势"):
        """绘制股票价格走势图"""
        fig, ax = plt.subplots(figsize=(14, 7))
        
        # 绘制收盘价和均线
        ax.plot(df.index, df['close'], label='收盘价', linewidth=2)
        if 'MA5' in df.columns:
            ax.plot(df.index, df['MA5'], label='MA5', linewidth=1)
        if 'MA20' in df.columns:
            ax.plot(df.index, df['MA20'], label='MA20', linewidth=1)
        if 'MA60' in df.columns:
            ax.plot(df.index, df['MA60'], label='MA60', linewidth=1)
        
        # 添加成交量子图
        ax2 = ax.twinx()
        ax2.bar(df.index, df['vol'], alpha=0.3, color='gray', label='成交量')
        ax2.set_ylabel('成交量')
        
        # 设置标题和标签
        ax.set_title(f"{ts_code} {title}", fontsize=16)
        ax.set_xlabel('日期')
        ax.set_ylabel('价格')
        ax.legend(loc='upper left')
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        ax.xaxis.set_major_formatter(date_format)
        fig.autofmt_xdate()
        
        plt.tight_layout()
        return fig
    
    def plot_technical_indicators(self, df, ts_code):
        """绘制技术指标图"""
        fig, axes = plt.subplots(3, 1, figsize=(14, 12), sharex=True)
        
        # 1. 价格和布林带
        axes[0].plot(df.index, df['close'], label='收盘价')
        axes[0].plot(df.index, df['BB_Upper'], 'g--', label='布林带上轨')
        axes[0].plot(df.index, df['BB_Middle'], 'y--', label='布林带中轨')
        axes[0].plot(df.index, df['BB_Lower'], 'r--', label='布林带下轨')
        axes[0].fill_between(df.index, df['BB_Lower'], df['BB_Upper'], alpha=0.1)
        axes[0].set_title(f"{ts_code} 价格与布林带")
        axes[0].legend()
        
        # 2. RSI指标
        axes[1].plot(df.index, df['RSI'])
        axes[1].axhline(y=70, color='r', linestyle='-', alpha=0.3)
        axes[1].axhline(y=30, color='g', linestyle='-', alpha=0.3)
        axes[1].set_title("RSI指标")
        axes[1].set_ylim(0, 100)
        
        # 3. MACD指标
        axes[2].plot(df.index, df['MACD'], label='MACD')
        axes[2].plot(df.index, df['MACD_Signal'], label='Signal')
        axes[2].bar(df.index, df['MACD_Hist'], label='Histogram')
        axes[2].set_title("MACD指标")
        axes[2].legend()
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        axes[2].xaxis.set_major_formatter(date_format)
        fig.autofmt_xdate()
        
        plt.tight_layout()
        return fig
    
    def plot_financial_indicators(self, df, ts_code):
        """绘制财务指标图"""
        if 'roe' not in df.columns or 'net_profit_margin' not in df.columns:
            print("数据中不包含ROE或净利润率指标")
            return None
        
        fig, axes = plt.subplots(2, 1, figsize=(14, 10), sharex=True)
        
        # 1. ROE和净利润率
        axes[0].plot(df.index, df['roe'], label='ROE(%)', color='blue')
        axes[0].set_title(f"{ts_code} ROE和净利润率")
        axes[0].set_ylabel('ROE(%)')
        axes[0].legend(loc='upper left')
        
        ax2 = axes[0].twinx()
        ax2.plot(df.index, df['net_profit_margin'], label='净利润率(%)', color='orange')
        ax2.set_ylabel('净利润率(%)')
        ax2.legend(loc='upper right')
        
        # 2. 资产负债率和流动比率
        if 'debt_to_assets' in df.columns and 'current_ratio' in df.columns:
            axes[1].plot(df.index, df['debt_to_assets'], label='资产负债率(%)', color='green')
            axes[1].set_title(f"{ts_code} 资产负债率和流动比率")
            axes[1].set_ylabel('资产负债率(%)')
            axes[1].legend(loc='upper left')
            
            ax3 = axes[1].twinx()
            ax3.plot(df.index, df['current_ratio'], label='流动比率', color='red')
            ax3.set_ylabel('流动比率')
            ax3.legend(loc='upper right')
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        axes[1].xaxis.set_major_formatter(date_format)
        fig.autofmt_xdate()
        
        plt.tight_layout()
        return fig
    
    def plot_correlation_matrix(self, df, title="特征相关性矩阵"):
        """绘制特征相关性矩阵"""
        # 选择数值列
        numeric_cols = df.select_dtypes(include=[np.number]).columns
        
        # 计算相关系数矩阵
        corr = df[numeric_cols].corr()
        
        # 创建掩码以只显示下三角
        mask = np.triu(np.ones_like(corr, dtype=bool))
        
        # 设置图形大小
        plt.figure(figsize=(14, 12))
        
        # 绘制热力图
        sns.heatmap(corr, mask=mask, annot=True, fmt=".2f", cmap='coolwarm', 
                    square=True, linewidths=.5, cbar_kws={"shrink": .8})
        
        plt.title(title, fontsize=16)
        plt.tight_layout()
        return plt.gcf()
    
    def plot_strategy_returns(self, results_dict, title="策略收益对比"):
        """绘制不同策略的收益曲线对比图"""
        plt.figure(figsize=(14, 8))
        
        for strategy_name, result_df in results_dict.items():
            plt.plot(result_df.index, result_df['cumulative_return'], 
                     label=strategy_name, linewidth=2)
        
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.title(title, fontsize=16)
        plt.xlabel('日期')
        plt.ylabel('累计收益率')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_strategy_performance(self, result_df, strategy_name):
        """绘制策略绩效综合图"""
        fig, axes = plt.subplots(2, 1, figsize=(14, 12), sharex=True)
        
        # 1. 权益曲线
        axes[0].plot(result_df.index, result_df['total'], label='总资产')
        axes[0].set_title(f"{strategy_name} 策略权益曲线", fontsize=14)
        axes[0].set_ylabel('资产价值(元)')
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)
        
        # 2. 每日收益率
        axes[1].bar(result_df.index, result_df['return'], alpha=0.6)
        axes[1].axhline(y=0, color='black', linestyle='-', alpha=0.3)
        axes[1].set_title(f"{strategy_name} 每日收益率", fontsize=14)
        axes[1].set_ylabel('收益率')
        axes[1].grid(True, alpha=0.3)
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        axes[1].xaxis.set_major_formatter(date_format)
        fig.autofmt_xdate()
        
        # 设置y轴为百分比格式
        axes[1].yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        plt.tight_layout()
        return fig
    
    def plot_drawdown(self, result_df, strategy_name):
        """绘制回撤曲线"""
        # 计算累计收益和最大回撤
        result_df['cumulative_return'] = (1 + result_df['return']).cumprod() - 1
        result_df['peak'] = result_df['cumulative_return'].cummax()
        result_df['drawdown'] = (result_df['cumulative_return'] - result_df['peak']) / (1 + result_df['peak'])
        
        plt.figure(figsize=(14, 7))
        
        # 绘制累计收益曲线
        plt.subplot(2, 1, 1)
        plt.plot(result_df.index, result_df['cumulative_return'], label='累计收益率')
        plt.plot(result_df.index, result_df['peak'], label='历史最高点', linestyle='--', alpha=0.7)
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.title(f"{strategy_name} 累计收益率与最大回撤", fontsize=14)
        plt.ylabel('收益率')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 绘制回撤曲线
        plt.subplot(2, 1, 2)
        plt.fill_between(result_df.index, result_df['drawdown'], 0, color='red', alpha=0.3)
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.5)
        plt.ylabel('回撤')
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        plt.gca().xaxis.set_major_formatter(date_format)
        plt.gcf().autofmt_xdate()
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_trades(self, price_df, signals_df, ts_code, strategy_name):
        """绘制交易信号图"""
        plt.figure(figsize=(14, 8))
        
        # 绘制收盘价
        plt.plot(price_df.index, price_df['close'], label='收盘价', linewidth=2)
        
        # 标记买入点
        buy_signals = signals_df[signals_df['signal'] == 1]
        plt.scatter(buy_signals.index, buy_signals['close'], marker='^', 
                   color='g', s=100, label='买入信号')
        
        # 标记卖出点
        sell_signals = signals_df[signals_df['signal'] == -1]
        plt.scatter(sell_signals.index, sell_signals['close'], marker='v', 
                   color='r', s=100, label='卖出信号')
        
        plt.title(f"{ts_code} {strategy_name} 交易信号图", fontsize=16)
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        plt.gca().xaxis.set_major_formatter(date_format)
        plt.gcf().autofmt_xdate()
        
        plt.tight_layout()
        return plt.gcf()
    
def get_stock_data(ts_code, start_date, end_date):
    """获取股票日线行情数据并进行简单处理"""
    try:
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        # 按交易日期升序排列
        df = df.sort_values('trade_date')
        df['trade_date'] = pd.to_datetime(df['trade_date'])
        df.set_index('trade_date', inplace=True)
        
        # 计算简单技术指标
        df['MA5'] = df['close'].rolling(window=5).mean()
        df['MA20'] = df['close'].rolling(window=20).mean()
        df['MA60'] = df['close'].rolling(window=60).mean()
        
        # 计算RSI指标
        delta = df['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        df['RSI'] = 100 - (100 / (1 + rs))
        
        # 计算布林带
        middle_band = df['close'].rolling(window=20).mean()
        std = df['close'].rolling(window=20).std()
        upper_band = middle_band + 2 * std
        lower_band = middle_band - 2 * std
        df['BB_Upper'] = upper_band
        df['BB_Middle'] = middle_band
        df['BB_Lower'] = lower_band
        
        # 计算MACD指标
        exp1 = df['close'].ewm(span=12, adjust=False).mean()
        exp2 = df['close'].ewm(span=26, adjust=False).mean()
        df['MACD'] = exp1 - exp2
        df['MACD_Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
        df['MACD_Hist'] = df['MACD'] - df['MACD_Signal']
        
        return df
    except Exception as e:
        print(f"获取{ts_code}日线数据出错: {e}")
        return None

def generate_sample_strategy_data(stock_data):
    """生成示例策略回测数据"""
    if stock_data is None:
        return None
    
    # 创建一个副本
    result = stock_data.copy()
    
    # 生成随机交易信号
    np.random.seed(42)  # 设置随机种子，确保结果可重现
    result['signal'] = 0
    
    # 随机生成买入信号(1)和卖出信号(-1)
    buy_dates = np.random.choice(result.index, size=10, replace=False)
    sell_dates = np.random.choice([d for d in result.index if d not in buy_dates], 
                                 size=10, replace=False)
    
    result.loc[buy_dates, 'signal'] = 1
    result.loc[sell_dates, 'signal'] = -1
    
    # 计算策略收益
    result['position'] = result['signal'].cumsum()  # 持仓状态
    result['position'].fillna(method='ffill', inplace=True)
    result['position'].fillna(0, inplace=True)
    
    # 计算策略每日收益
    result['strategy_return'] = result['position'].shift(1) * result['close'].pct_change()
    
    # 计算累计收益
    result['cumulative_return'] = (1 + result['strategy_return']).cumprod() - 1
    
    # 计算总资产（假设初始资金100000元）
    initial_capital = 100000
    result['total'] = initial_capital * (1 + result['cumulative_return'])
    
    return result

# 主函数：演示如何使用可视化工具
def main():
    # 获取股票数据
    stock_data = get_stock_data(ts_code='000001.SZ', start_date='20200101', end_date='20231231')
    
    if stock_data is None:
        print("无法获取股票数据，请检查Tushare token和网络连接")
        return
    
    # 生成示例策略数据
    strategy_data = generate_sample_strategy_data(stock_data)
    
    # 创建可视化工具实例
    visualizer = QuantVisualizer()
    
    # 1. 绘制股票价格和均线
    price_fig = visualizer.plot_stock_price(stock_data, '000001.SZ', "平安银行价格走势")
    price_fig.savefig('price_chart.png')
    
    # 2. 绘制技术指标
    tech_fig = visualizer.plot_technical_indicators(stock_data, '000001.SZ')
    tech_fig.savefig('technical_indicators.png')
    
    # 3. 绘制交易信号图
    signals_fig = visualizer.plot_trades(stock_data, strategy_data, '000001.SZ', "示例策略")
    signals_fig.savefig('trading_signals.png')
    
    # 4. 绘制策略绩效
    performance_fig = visualizer.plot_strategy_performance(strategy_data, "示例策略")
    performance_fig.savefig('strategy_performance.png')
    
    # 5. 绘制回撤曲线
    drawdown_fig = visualizer.plot_drawdown(strategy_data, "示例策略")
    drawdown_fig.savefig('strategy_drawdown.png')
    
    print("可视化图表已成功生成并保存！")

if __name__ == "__main__":
    main()


# In[12]:


def process_financial_data(daily_data, financial_data):
    """处理财务数据并与日线数据对齐"""
    # 假设financial_data包含以下列：
    # trade_date, profit_grow_rate, debt_to_assets, roe, gross_margin
    
    # 将财务数据按日期升序排列
    financial_data = financial_data.sort_values('trade_date')
    
    # 使用前向填充将财务数据对齐到每日
    merged_data = pd.merge_asof(
        daily_data, 
        financial_data[['trade_date', 'profit_grow_rate', 'debt_to_assets', 'roe', 'gross_margin']],
        on='trade_date',
        direction='backward'
    )
    
    # 处理缺失值
    merged_data = merged_data.fillna(method='ffill')
    
    # 处理异常值
    for col in ['profit_grow_rate', 'debt_to_assets', 'roe', 'gross_margin']:
        # 截断极端值（上下1%分位）
        lower = merged_data[col].quantile(0.01)
        upper = merged_data[col].quantile(0.99)
        merged_data[col] = merged_data[col].clip(lower, upper)
    
    return merged_data


# In[13]:


import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
from matplotlib.dates import DateFormatter
import matplotlib.ticker as mticker
from sklearn.metrics import confusion_matrix, classification_report, roc_curve, auc

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
sns.set_style('whitegrid')

class StrategyVisualizer:
    """量化策略可视化工具类"""
    
    def __init__(self):
        pass
    
    def plot_strategy_signals(self, df, ts_code, strategy_name):
        """绘制策略信号图"""
        plt.figure(figsize=(14, 8))
        
        # 绘制收盘价
        plt.plot(df.index, df['close'], label='收盘价', linewidth=2)
        
        # 标记买入点
        buy_points = df[df['signal'] == 1]
        plt.scatter(buy_points.index, buy_points['close'], marker='^', 
                   color='g', s=100, label='买入信号')
        
        # 标记卖出点
        sell_points = df[df['signal'] == -1]
        plt.scatter(sell_points.index, sell_points['close'], marker='v', 
                   color='r', s=100, label='卖出信号')
        
        plt.title(f"{ts_code} {strategy_name} 交易信号图", fontsize=16)
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        plt.gca().xaxis.set_major_formatter(date_format)
        plt.gcf().autofmt_xdate()
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_cumulative_returns(self, strategies_dict, title="策略累计收益率对比"):
        """绘制不同策略的累计收益率对比图"""
        plt.figure(figsize=(14, 8))
        
        for strategy_name, df in strategies_dict.items():
            # 计算累计收益率
            df['cumulative_return'] = (1 + df['daily_return']).cumprod() - 1
            plt.plot(df.index, df['cumulative_return'], label=strategy_name, linewidth=2)
        
        # 添加基准（持有策略）
        plt.plot(strategies_dict[list(strategies_dict.keys())[0]].index, 
                 (1 + strategies_dict[list(strategies_dict.keys())[0]]['benchmark_return']).cumprod() - 1, 
                 label='持有策略', linestyle='--', color='gray')
        
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.title(title, fontsize=16)
        plt.xlabel('日期')
        plt.ylabel('累计收益率')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_drawdown(self, df, strategy_name):
        """绘制回撤曲线"""
        # 计算累计收益和最大回撤
        df['cumulative_return'] = (1 + df['daily_return']).cumprod() - 1
        df['peak'] = df['cumulative_return'].cummax()
        df['drawdown'] = (df['cumulative_return'] - df['peak']) / (1 + df['peak'])
        
        plt.figure(figsize=(14, 7))
        
        # 绘制累计收益曲线
        plt.subplot(2, 1, 1)
        plt.plot(df.index, df['cumulative_return'], label='累计收益率')
        plt.plot(df.index, df['peak'], label='历史最高点', linestyle='--', alpha=0.7)
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.title(f"{strategy_name} 累计收益率与最大回撤", fontsize=14)
        plt.ylabel('收益率')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 绘制回撤曲线
        plt.subplot(2, 1, 2)
        plt.fill_between(df.index, df['drawdown'], 0, color='red', alpha=0.3)
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.5)
        plt.ylabel('回撤')
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        plt.gca().xaxis.set_major_formatter(date_format)
        plt.gcf().autofmt_xdate()
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_monthly_returns(self, df, strategy_name):
        """绘制月度收益率热力图"""
        # 计算月度收益率
        monthly_returns = df['daily_return'].resample('M').apply(lambda x: (1 + x).prod() - 1)
        
        # 创建年月表格
        monthly_df = monthly_returns.reset_index()
        monthly_df['year'] = monthly_df['trade_date'].dt.year
        monthly_df['month'] = monthly_df['trade_date'].dt.month
        monthly_df = monthly_df.pivot(index='year', columns='month', values='daily_return')
        
        # 绘制热力图
        plt.figure(figsize=(14, 8))
        cmap = sns.diverging_palette(10, 130, as_cmap=True)
        ax = sns.heatmap(monthly_df, annot=True, fmt=".2%", cmap=cmap, 
                         center=0, linewidths=.5, cbar_kws={"shrink": .8})
        
        plt.title(f"{strategy_name} 月度收益率热力图", fontsize=16)
        plt.tight_layout()
        return plt.gcf()
    
    def plot_performance_metrics(self, strategies_dict):
        """绘制策略绩效指标对比图"""
        # 计算各策略的绩效指标
        metrics = []
        
        for strategy_name, df in strategies_dict.items():
            # 计算累计收益率
            cumulative_return = (1 + df['daily_return']).cumprod().iloc[-1] - 1
            
            # 计算年化收益率
            annual_return = (1 + cumulative_return) ** (252 / len(df)) - 1
            
            # 计算夏普比率
            risk_free_rate = 0.03  # 假设无风险利率为3%
            sharpe_ratio = (annual_return - risk_free_rate) / (df['daily_return'].std() * np.sqrt(252))
            
            # 计算最大回撤
            df['cumulative'] = (1 + df['daily_return']).cumprod()
            df['peak'] = df['cumulative'].cummax()
            df['drawdown'] = (df['cumulative'] - df['peak']) / df['peak']
            max_drawdown = df['drawdown'].min()
            
            # 计算胜率
            win_rate = len(df[df['daily_return'] > 0]) / len(df)
            
            metrics.append({
                '策略': strategy_name,
                '年化收益率': annual_return,
                '夏普比率': sharpe_ratio,
                '最大回撤': max_drawdown,
                '胜率': win_rate
            })
        
        # 转换为DataFrame
        metrics_df = pd.DataFrame(metrics)
        
        # 绘制绩效对比图
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        
        # 年化收益率
        sns.barplot(x='策略', y='年化收益率', data=metrics_df, ax=axes[0, 0])
        axes[0, 0].set_title('年化收益率对比')
        axes[0, 0].yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 夏普比率
        sns.barplot(x='策略', y='夏普比率', data=metrics_df, ax=axes[0, 1])
        axes[0, 1].set_title('夏普比率对比')
        
        # 最大回撤
        sns.barplot(x='策略', y='最大回撤', data=metrics_df, ax=axes[1, 0])
        axes[1, 0].set_title('最大回撤对比')
        axes[1, 0].yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 胜率
        sns.barplot(x='策略', y='胜率', data=metrics_df, ax=axes[1, 1])
        axes[1, 1].set_title('胜率对比')
        axes[1, 1].yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        plt.tight_layout()
        return fig
    
    def plot_ml_feature_importance(self, model, feature_names, top_n=10):
        """绘制机器学习特征重要性图"""
        if hasattr(model, 'feature_importances_'):
            # 获取特征重要性
            importances = model.feature_importances_
            indices = np.argsort(importances)[-top_n:]
            
            # 绘制特征重要性图
            plt.figure(figsize=(10, 6))
            plt.title('特征重要性排名')
            plt.barh(range(len(indices)), importances[indices], align='center')
            plt.yticks(range(len(indices)), [feature_names[i] for i in indices])
            plt.xlabel('重要性')
            plt.tight_layout()
            return plt.gcf()
        else:
            print("模型不支持特征重要性计算")
            return None
    
    def plot_ml_confusion_matrix(self, y_true, y_pred, classes=['下跌', '上涨']):
        """绘制机器学习混淆矩阵"""
        # 计算混淆矩阵
        cm = confusion_matrix(y_true, y_pred)
        
        # 绘制混淆矩阵
        plt.figure(figsize=(8, 6))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                   xticklabels=classes, yticklabels=classes)
        plt.title('混淆矩阵')
        plt.xlabel('预测类别')
        plt.ylabel('实际类别')
        plt.tight_layout()
        return plt.gcf()
    
    def plot_ml_roc_curve(self, y_true, y_score):
        """绘制机器学习ROC曲线"""
        # 计算ROC曲线
        fpr, tpr, _ = roc_curve(y_true, y_score)
        roc_auc = auc(fpr, tpr)
        
        # 绘制ROC曲线
        plt.figure(figsize=(8, 6))
        plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc:.2f})')
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('假阳性率')
        plt.ylabel('真阳性率')
        plt.title('ROC曲线')
        plt.legend(loc="lower right")
        plt.tight_layout()
        return plt.gcf()


# In[16]:


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.dates import DateFormatter
import matplotlib.ticker as mticker
import tushare as ts
from sklearn.ensemble import RandomForestClassifier

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
sns.set_style('whitegrid')

# 设置Tushare token（需替换为个人token）
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

class StrategyVisualizer:
    """量化策略可视化工具类"""
    
    def __init__(self):
        pass
    
    def plot_strategy_signals(self, df, ts_code, strategy_name):
        """绘制策略信号图"""
        plt.figure(figsize=(14, 8))
        
        # 绘制收盘价
        plt.plot(df.index, df['close'], label='收盘价', linewidth=2)
        
        # 标记买入点
        buy_points = df[df['signal'] == 1]
        plt.scatter(buy_points.index, buy_points['close'], marker='^', 
                   color='g', s=100, label='买入信号')
        
        # 标记卖出点
        sell_points = df[df['signal'] == -1]
        plt.scatter(sell_points.index, sell_points['close'], marker='v', 
                   color='r', s=100, label='卖出信号')
        
        plt.title(f"{ts_code} {strategy_name} 交易信号图", fontsize=16)
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        plt.gca().xaxis.set_major_formatter(date_format)
        plt.gcf().autofmt_xdate()
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_cumulative_returns(self, strategies_dict, title="策略累计收益率对比"):
        """绘制不同策略的累计收益率对比图"""
        plt.figure(figsize=(14, 8))
        
        for strategy_name, df in strategies_dict.items():
            # 计算累计收益率
            df['cumulative_return'] = (1 + df['daily_return']).cumprod() - 1
            plt.plot(df.index, df['cumulative_return'], label=strategy_name, linewidth=2)
        
        # 添加基准（持有策略）
        plt.plot(strategies_dict[list(strategies_dict.keys())[0]].index, 
                 (1 + strategies_dict[list(strategies_dict.keys())[0]]['benchmark_return']).cumprod() - 1, 
                 label='持有策略', linestyle='--', color='gray')
        
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.title(title, fontsize=16)
        plt.xlabel('日期')
        plt.ylabel('累计收益率')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_drawdown(self, df, strategy_name):
        """绘制回撤曲线"""
        # 计算累计收益和最大回撤
        df['cumulative_return'] = (1 + df['daily_return']).cumprod() - 1
        df['peak'] = df['cumulative_return'].cummax()
        df['drawdown'] = (df['cumulative_return'] - df['peak']) / (1 + df['peak'])
        
        plt.figure(figsize=(14, 7))
        
        # 绘制累计收益曲线
        plt.subplot(2, 1, 1)
        plt.plot(df.index, df['cumulative_return'], label='累计收益率')
        plt.plot(df.index, df['peak'], label='历史最高点', linestyle='--', alpha=0.7)
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.title(f"{strategy_name} 累计收益率与最大回撤", fontsize=14)
        plt.ylabel('收益率')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 绘制回撤曲线
        plt.subplot(2, 1, 2)
        plt.fill_between(df.index, df['drawdown'], 0, color='red', alpha=0.3)
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.5)
        plt.ylabel('回撤')
        plt.grid(True, alpha=0.3)
        
        # 设置y轴为百分比格式
        plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        plt.gca().xaxis.set_major_formatter(date_format)
        plt.gcf().autofmt_xdate()
        
        plt.tight_layout()
        return plt.gcf()
    
    def plot_monthly_returns(self, df, strategy_name):
        """绘制月度收益率热力图"""
        # 计算月度收益率
        monthly_returns = df['daily_return'].resample('M').apply(lambda x: (1 + x).prod() - 1)
        
        # 创建年月表格
        monthly_df = monthly_returns.reset_index()
        monthly_df['year'] = monthly_df['trade_date'].dt.year
        monthly_df['month'] = monthly_df['trade_date'].dt.month
        monthly_df = monthly_df.pivot(index='year', columns='month', values='daily_return')
        
        # 绘制热力图
        plt.figure(figsize=(14, 8))
        cmap = sns.diverging_palette(10, 130, as_cmap=True)
        ax = sns.heatmap(monthly_df, annot=True, fmt=".2%", cmap=cmap, 
                         center=0, linewidths=.5, cbar_kws={"shrink": .8})
        
        plt.title(f"{strategy_name} 月度收益率热力图", fontsize=16)
        plt.tight_layout()
        return plt.gcf()
    
    def plot_performance_metrics(self, strategies_dict):
        """绘制策略绩效指标对比图"""
        # 计算各策略的绩效指标
        metrics = []
        
        for strategy_name, df in strategies_dict.items():
            # 计算累计收益率
            cumulative_return = (1 + df['daily_return']).cumprod().iloc[-1] - 1
            
            # 计算年化收益率
            annual_return = (1 + cumulative_return) ** (252 / len(df)) - 1
            
            # 计算夏普比率
            risk_free_rate = 0.03  # 假设无风险利率为3%
            sharpe_ratio = (annual_return - risk_free_rate) / (df['daily_return'].std() * np.sqrt(252))
            
            # 计算最大回撤
            df['cumulative'] = (1 + df['daily_return']).cumprod()
            df['peak'] = df['cumulative'].cummax()
            df['drawdown'] = (df['cumulative'] - df['peak']) / df['peak']
            max_drawdown = df['drawdown'].min()
            
            # 计算胜率
            win_rate = len(df[df['daily_return'] > 0]) / len(df)
            
            metrics.append({
                '策略': strategy_name,
                '年化收益率': annual_return,
                '夏普比率': sharpe_ratio,
                '最大回撤': max_drawdown,
                '胜率': win_rate
            })
        
        # 转换为DataFrame
        metrics_df = pd.DataFrame(metrics)
        
        # 绘制绩效对比图
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        
        # 年化收益率
        sns.barplot(x='策略', y='年化收益率', data=metrics_df, ax=axes[0, 0])
        axes[0, 0].set_title('年化收益率对比')
        axes[0, 0].yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 夏普比率
        sns.barplot(x='策略', y='夏普比率', data=metrics_df, ax=axes[0, 1])
        axes[0, 1].set_title('夏普比率对比')
        
        # 最大回撤
        sns.barplot(x='策略', y='最大回撤', data=metrics_df, ax=axes[1, 0])
        axes[1, 0].set_title('最大回撤对比')
        axes[1, 0].yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        # 胜率
        sns.barplot(x='策略', y='胜率', data=metrics_df, ax=axes[1, 1])
        axes[1, 1].set_title('胜率对比')
        axes[1, 1].yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
        
        plt.tight_layout()
        return fig
    
    def plot_ml_feature_importance(self, model, feature_names, top_n=10):
        """绘制机器学习特征重要性图"""
        if hasattr(model, 'feature_importances_'):
            # 获取特征重要性
            importances = model.feature_importances_
            indices = np.argsort(importances)[-top_n:]
            
            # 绘制特征重要性图
            plt.figure(figsize=(10, 6))
            plt.title('特征重要性排名')
            plt.barh(range(len(indices)), importances[indices], align='center')
            plt.yticks(range(len(indices)), [feature_names[i] for i in indices])
            plt.xlabel('重要性')
            plt.tight_layout()
            return plt.gcf()
        else:
            print("模型不支持特征重要性计算")
            return None
    
    def plot_ml_confusion_matrix(self, y_true, y_pred, classes=['下跌', '上涨']):
        """绘制机器学习混淆矩阵"""
        # 计算混淆矩阵
        cm = confusion_matrix(y_true, y_pred)
        
        # 绘制混淆矩阵
        plt.figure(figsize=(8, 6))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                   xticklabels=classes, yticklabels=classes)
        plt.title('混淆矩阵')
        plt.xlabel('预测类别')
        plt.ylabel('实际类别')
        plt.tight_layout()
        return plt.gcf()
    
    def plot_ml_roc_curve(self, y_true, y_score):
        """绘制机器学习ROC曲线"""
        # 计算ROC曲线
        fpr, tpr, _ = roc_curve(y_true, y_score)
        roc_auc = auc(fpr, tpr)
        
        # 绘制ROC曲线
        plt.figure(figsize=(8, 6))
        plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc:.2f})')
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('假阳性率')
        plt.ylabel('真阳性率')
        plt.title('ROC曲线')
        plt.legend(loc="lower right")
        plt.tight_layout()
        return plt.gcf()

def get_stock_data(ts_code, start_date, end_date):
    """获取股票日线行情数据并进行简单处理"""
    try:
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        # 按交易日期升序排列
        df = df.sort_values('trade_date')
        df['trade_date'] = pd.to_datetime(df['trade_date'])
        df.set_index('trade_date', inplace=True)
        
        # 计算简单技术指标
        df['MA5'] = df['close'].rolling(window=5).mean()
        df['MA20'] = df['close'].rolling(window=20).mean()
        df['MA60'] = df['close'].rolling(window=60).mean()
        
        # 计算RSI指标
        delta = df['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        df['RSI'] = 100 - (100 / (1 + rs))
        
        # 计算MACD指标
        exp1 = df['close'].ewm(span=12, adjust=False).mean()
        exp2 = df['close'].ewm(span=26, adjust=False).mean()
        df['DIF'] = exp1 - exp2
        df['DEA'] = df['DIF'].ewm(span=9, adjust=False).mean()
        df['MACD_Hist'] = df['DIF'] - df['DEA']
        
        # 计算KDJ指标
        low_9 = df['low'].rolling(window=9).min()
        high_9 = df['high'].rolling(window=9).max()
        df['K'] = 100 * ((df['close'] - low_9) / (high_9 - low_9))
        df['D'] = df['K'].rolling(window=3).mean()
        df['J'] = 3 * df['K'] - 2 * df['D']
        
        # 计算基准收益（持有策略）
        df['benchmark_return'] = df['close'].pct_change()
        
        return df
    except Exception as e:
        print(f"获取{ts_code}日线数据出错: {e}")
        return None

def generate_technical_signals(df):
    """基于量价分析的传统量化策略 - 信号生成"""
    # 确保数据包含所需列
    required_columns = ['close', 'MA5', 'MA20', 'DIF', 'DEA', 'K', 'D']
    for col in required_columns:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要的列: {col}")
    
    # 初始化信号列
    df['signal'] = 0
    
    # 买入条件
    buy_condition = (df['MA5'] > df['MA20']) & (df['MA5'].shift(1) <= df['MA20'].shift(1)) &                     (df['DIF'] > df['DEA']) & (df['DIF'].shift(1) <= df['DEA'].shift(1)) &                     (df['K'] > df['D']) & (df['K'].shift(1) <= df['D'].shift(1))
    
    # 卖出条件
    sell_condition = ((df['MA5'] < df['MA20']) & (df['MA5'].shift(1) >= df['MA20'].shift(1))) |                      ((df['DIF'] < df['DEA']) & (df['DIF'].shift(1) >= df['DEA'].shift(1))) |                      ((df['K'] < df['D']) & (df['K'].shift(1) >= df['D'].shift(1)))
    
    # 设置信号
    df.loc[buy_condition, 'signal'] = 1
    df.loc[sell_condition, 'signal'] = -1
    
    return df

def generate_fundamental_signals(df, debt_ratio_threshold=0.6):
    """基于财务指标的量化策略 - 信号生成"""
    # 确保数据包含所需列
    required_columns = ['close', 'MA5', 'MA20', 'profit_grow_rate', 'debt_to_assets']
    for col in required_columns:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要的列: {col}")
    
    # 初始化信号列
    df['signal'] = 0
    
    # 定义资产负债率适中范围
    moderate_debt = (df['debt_to_assets'] > 0.2) & (df['debt_to_assets'] <= debt_ratio_threshold)
    high_debt = df['debt_to_assets'] > debt_ratio_threshold
    
    # 买入条件
    buy_condition = (df['profit_grow_rate'] > 0) & moderate_debt &                     (df['MA5'] > df['MA20']) & (df['MA5'].shift(1) <= df['MA20'].shift(1))
    
    # 卖出条件
    sell_condition = (df['profit_grow_rate'] <= 0) | high_debt |                     ((df['MA5'] < df['MA20']) & (df['MA5'].shift(1) >= df['MA20'].shift(1)))
    
    # 设置信号
    df.loc[buy_condition, 'signal'] = 1
    df.loc[sell_condition, 'signal'] = -1
    
    return df

def create_machine_learning_dataset(df, features, target_col='future_return', days_ahead=5):
    """创建机器学习数据集"""
    # 确保数据包含所需特征列
    for feature in features:
        if feature not in df.columns:
            raise ValueError(f"数据缺少特征列: {feature}")
    
    # 计算目标变量：未来n天的收益率
    df[target_col] = df['close'].pct_change(days_ahead).shift(-days_ahead)
    
    # 创建标签：1表示正收益，0表示负收益
    df['label'] = (df[target_col] > 0).astype(int)
    
    # 移除缺失值
    df = df.dropna()
    
    # 划分特征和标签
    X = df[features]
    y = df['label']
    
    # 划分训练集和测试集（70%训练，30%测试）
    split_idx = int(len(df) * 0.7)
    X_train, X_test = X.iloc[:split_idx], X.iloc[split_idx:]
    y_train, y_test = y.iloc[:split_idx], y.iloc[split_idx:]
    
    # 保存日期索引用于后续分析
    train_dates = df.index[:split_idx]
    test_dates = df.index[split_idx:]
    
    return X_train, X_test, y_train, y_test, train_dates, test_dates, df

def train_random_forest_model(X_train, y_train, n_estimators=100, random_state=42):
    """训练随机森林模型"""
    model = RandomForestClassifier(
        n_estimators=n_estimators,
        random_state=random_state,
        n_jobs=-1  # 使用所有CPU核心
    )
    
    # 训练模型
    model.fit(X_train, y_train)
    
    print(f"训练集准确率: {model.score(X_train, y_train):.4f}")
    
    return model

def generate_ml_signals(df, model, features, threshold=0.6):
    """使用机器学习模型生成交易信号"""
    # 预测概率
    df['pred_prob'] = model.predict_proba(df[features])[:, 1]
    
    # 初始化信号列
    df['signal'] = 0
    
    # 生成信号：当预测上涨概率大于阈值时买入，小于(1-阈值)时卖出
    df.loc[df['pred_prob'] >= threshold, 'signal'] = 1
    df.loc[df['pred_prob'] <= (1 - threshold), 'signal'] = -1
    
    return df

def backtest_strategy(df, initial_capital=100000):
    """回测策略"""
    # 复制数据
    result = df.copy()
    
    # 初始化持仓和现金
    result['position'] = 0  # 持仓数量
    result['cash'] = initial_capital  # 现金余额
    result['total'] = initial_capital  # 总资产
    
    position = 0
    cash = initial_capital
    
    # 回测逻辑
    for i in range(1, len(result)):
        # 获取前一天的信号
        signal = result.iloc[i-1]['signal']
        
        # 获取当天价格
        price = result.iloc[i]['close']
        
        # 执行交易
        if signal == 1 and position == 0:  # 买入
            position = cash // price  # 买入尽可能多的股票
            cash = cash - position * price
        elif signal == -1 and position > 0:  # 卖出
            cash = cash + position * price
            position = 0
        
        # 更新持仓和总资产
        result.iloc[i, result.columns.get_loc('position')] = position
        result.iloc[i, result.columns.get_loc('cash')] = cash
        result.iloc[i, result.columns.get_loc('total')] = cash + position * price
    
    # 计算每日收益率
    result['daily_return'] = result['total'].pct_change()
    
    return result

# 主函数：演示如何使用可视化工具
def main():
    # 获取股票数据
    stock_data = get_stock_data(ts_code='000001.SZ', start_date='20200101', end_date='20231231')
    
    if stock_data is None:
        print("无法获取股票数据，请检查Tushare token和网络连接")
        return
    
    # 为演示目的，添加模拟财务数据
    stock_data['profit_grow_rate'] = np.random.normal(0.1, 0.2, len(stock_data))
    stock_data['debt_to_assets'] = np.random.normal(0.4, 0.1, len(stock_data))
    
    # 1. 量价策略回测
    tech_strategy_df = generate_technical_signals(stock_data.copy())
    tech_strategy_df = backtest_strategy(tech_strategy_df)
    
    # 2. 财务策略回测
    fundamental_strategy_df = generate_fundamental_signals(stock_data.copy())
    fundamental_strategy_df = backtest_strategy(fundamental_strategy_df)
    
    # 3. 机器学习策略回测
    # 定义特征
    features = ['MA5', 'MA20', 'MA60', 'RSI', 'DIF', 'DEA', 'K', 'D', 'J', 
                'profit_grow_rate', 'debt_to_assets']
    
    # 创建数据集
    X_train, X_test, y_train, y_test, train_dates, test_dates, ml_data = create_machine_learning_dataset(
        stock_data.copy(), features)
    
    # 训练模型
    model = train_random_forest_model(X_train, y_train)
    
    # 生成预测信号
    ml_data = generate_ml_signals(ml_data, model, features)
    
    # 只保留测试集部分进行回测
    ml_test_data = ml_data.loc[test_dates]
    ml_strategy_df = backtest_strategy(ml_test_data)
    
    # 创建可视化工具实例
    visualizer = StrategyVisualizer()
    
    # 1. 绘制各策略的交易信号图
    tech_signals_fig = visualizer.plot_strategy_signals(tech_strategy_df, '000001.SZ', "量价策略")
    tech_signals_fig.savefig('tech_strategy_signals.png')
    
    fundamental_signals_fig = visualizer.plot_strategy_signals(fundamental_strategy_df, '000001.SZ', "财务策略")
    fundamental_signals_fig.savefig('fundamental_strategy_signals.png')
    
    ml_signals_fig = visualizer.plot_strategy_signals(ml_strategy_df, '000001.SZ', "机器学习策略")
    ml_signals_fig.savefig('ml_strategy_signals.png')
    
    # 2. 绘制累计收益率对比图
    strategies_dict = {
        "量价策略": tech_strategy_df,
        "财务策略": fundamental_strategy_df,
        "机器学习策略": ml_strategy_df
    }
    returns_fig = visualizer.plot_cumulative_returns(strategies_dict)
    returns_fig.savefig('strategy_returns_comparison.png')
    
    # 3. 绘制各策略的回撤曲线
    tech_drawdown_fig = visualizer.plot_drawdown(tech_strategy_df, "量价策略")
    tech_drawdown_fig.savefig('tech_strategy_drawdown.png')
    
    fundamental_drawdown_fig = visualizer.plot_drawdown(fundamental_strategy_df, "财务策略")
    fundamental_drawdown_fig.savefig('fundamental_strategy_drawdown.png')
    
    ml_drawdown_fig = visualizer.plot_drawdown(ml_strategy_df, "机器学习策略")
    ml_drawdown_fig.savefig('ml_strategy_drawdown.png')
    
    # 4. 绘制月度收益率热力图
    tech_monthly_fig = visualizer.plot_monthly_returns(tech_strategy_df, "量价策略")
    tech_monthly_fig.savefig('tech_monthly_returns.png')
    
    # 5. 绘制策略绩效指标对比图
    metrics_fig = visualizer.plot_performance_metrics(strategies_dict)
    metrics_fig.savefig('strategy_metrics_comparison.png')
    
    # 6. 针对机器学习策略，绘制特征重要性图
    feature_importance_fig = visualizer.plot_ml_feature_importance(
        model, X_train.columns, top_n=10)
    if feature_importance_fig:
        feature_importance_fig.savefig('ml_feature_importance.png')
    
    print("可视化图表已成功生成并保存！")

if __name__ == "__main__":
    main()


# In[18]:


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.dates import DateFormatter
import matplotlib.ticker as mticker
import tushare as ts
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import confusion_matrix, classification_report, roc_curve, auc
import logging
from datetime import datetime, timedelta

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
sns.set_style('whitegrid')

class Strategy:
    """策略基类"""
    
    def __init__(self, name):
        self.name = name
        self.signals = pd.DataFrame()
        
    def generate_signals(self, data):
        """生成交易信号，子类需实现此方法"""
        raise NotImplementedError("子类必须实现generate_signals方法")
        
    def get_signals(self):
        """获取生成的信号"""
        return self.signals

class TechnicalStrategy(Strategy):
    """基于量价分析的传统量化策略"""
    
    def __init__(self):
        super().__init__("量价策略")
        
    def generate_signals(self, data):
        """生成量价策略信号"""
        # 确保数据包含所需列
        required_columns = ['close', 'MA5', 'MA20', 'DIF', 'DEA', 'K', 'D']
        for col in required_columns:
            if col not in data.columns:
                raise ValueError(f"数据缺少必要的列: {col}")
        
        # 初始化信号列
        df = data.copy()
        df['signal'] = 0
        
        # 买入条件：MA5上穿MA20且MACD金叉且KDJ金叉
        buy_condition = (df['MA5'] > df['MA20']) & (df['MA5'].shift(1) <= df['MA20'].shift(1)) &                         (df['DIF'] > df['DEA']) & (df['DIF'].shift(1) <= df['DEA'].shift(1)) &                         (df['K'] > df['D']) & (df['K'].shift(1) <= df['D'].shift(1))
        
        # 卖出条件：MA5下穿MA20或MACD死叉或KDJ死叉
        sell_condition = ((df['MA5'] < df['MA20']) & (df['MA5'].shift(1) >= df['MA20'].shift(1))) |                          ((df['DIF'] < df['DEA']) & (df['DIF'].shift(1) >= df['DEA'].shift(1))) |                          ((df['K'] < df['D']) & (df['K'].shift(1) >= df['D'].shift(1)))
        
        # 设置信号
        df.loc[buy_condition, 'signal'] = 1
        df.loc[sell_condition, 'signal'] = -1
        
        self.signals = df[['signal']]
        return self.signals

class FundamentalStrategy(Strategy):
    """基于财务指标的量化策略"""
    
    def __init__(self, debt_ratio_threshold=0.6):
        super().__init__("财务策略")
        self.debt_ratio_threshold = debt_ratio_threshold
        
    def generate_signals(self, data):
        """生成财务策略信号"""
        # 确保数据包含所需列
        required_columns = ['close', 'MA5', 'MA20', 'profit_grow_rate', 'debt_to_assets']
        for col in required_columns:
            if col not in data.columns:
                raise ValueError(f"数据缺少必要的列: {col}")
        
        # 初始化信号列
        df = data.copy()
        df['signal'] = 0
        
        # 定义资产负债率适中范围
        moderate_debt = (df['debt_to_assets'] > 0.2) & (df['debt_to_assets'] <= self.debt_ratio_threshold)
        high_debt = df['debt_to_assets'] > self.debt_ratio_threshold
        
        # 买入条件：净利润增长率为正且资产负债率适中且MA5上穿MA20
        buy_condition = (df['profit_grow_rate'] > 0) & moderate_debt &                         (df['MA5'] > df['MA20']) & (df['MA5'].shift(1) <= df['MA20'].shift(1))
        
        # 卖出条件：净利润增长率为负或资产负债率过高或MA5下穿MA20
        sell_condition = (df['profit_grow_rate'] <= 0) | high_debt |                         ((df['MA5'] < df['MA20']) & (df['MA5'].shift(1) >= df['MA20'].shift(1)))
        
        # 设置信号
        df.loc[buy_condition, 'signal'] = 1
        df.loc[sell_condition, 'signal'] = -1
        
        self.signals = df[['signal']]
        return self.signals

class MLStrategy(Strategy):
    """基于机器学习的量化策略"""
    
    def __init__(self, features, model=None, threshold=0.6):
        super().__init__("机器学习策略")
        self.features = features
        self.model = model
        self.threshold = threshold
        
    def train_model(self, X_train, y_train):
        """训练机器学习模型"""
        if self.model is None:
            # 使用随机森林作为默认模型
            self.model = RandomForestClassifier(n_estimators=100, random_state=42, n_jobs=-1)
            
        # 训练模型
        self.model.fit(X_train, y_train)
        logger.info(f"模型训练完成，训练集准确率: {self.model.score(X_train, y_train):.4f}")
        
        return self.model
    
    def generate_signals(self, data):
        """生成机器学习策略信号"""
        if self.model is None:
            raise ValueError("模型未训练，请先调用train_model方法")
        
        # 确保数据包含所需特征列
        for feature in self.features:
            if feature not in data.columns:
                raise ValueError(f"数据缺少特征列: {feature}")
        
        # 预测概率
        df = data.copy()
        df['pred_prob'] = self.model.predict_proba(df[self.features])[:, 1]
        
        # 初始化信号列
        df['signal'] = 0
        
        # 生成信号：当预测上涨概率大于阈值时买入，小于(1-阈值)时卖出
        df.loc[df['pred_prob'] >= self.threshold, 'signal'] = 1
        df.loc[df['pred_prob'] <= (1 - self.threshold), 'signal'] = -1
        
        self.signals = df[['signal']]
        return self.signals

class StopLossTakeProfit:
    """止损止盈机制"""
    
    def __init__(self, stop_loss_pct=0.05, take_profit_pct=0.15, partial_take_profit_pct=0.5):
        """
        初始化止损止盈参数
        
        参数:
            stop_loss_pct: 止损百分比，默认为5%
            take_profit_pct: 止盈百分比，默认为15%
            partial_take_profit_pct: 部分止盈比例，默认为50%
        """
        self.stop_loss_pct = stop_loss_pct
        self.take_profit_pct = take_profit_pct
        self.partial_take_profit_pct = partial_take_profit_pct
        self.entry_price = 0  # 入场价格
        self.position = 0  # 当前持仓量
        
    def update_position(self, position, price):
        """更新持仓信息"""
        if position != self.position:
            self.position = position
            if position > 0:
                self.entry_price = price  # 记录入场价格
                logger.info(f"更新持仓: {position}，入场价格: {price}")
    
    def check_stop_loss(self, current_price):
        """检查是否触发止损条件"""
        if self.position <= 0:
            return False
            
        # 计算当前亏损比例
        loss_pct = (self.entry_price - current_price) / self.entry_price
        
        if loss_pct >= self.stop_loss_pct:
            logger.info(f"触发止损: 入场价={self.entry_price}, 当前价={current_price}, 亏损比例={loss_pct:.2%}")
            return True
            
        return False
    
    def check_take_profit(self, current_price):
        """检查是否触发止盈条件"""
        if self.position <= 0:
            return False
            
        # 计算当前盈利比例
        profit_pct = (current_price - self.entry_price) / self.entry_price
        
        if profit_pct >= self.take_profit_pct:
            logger.info(f"触发止盈: 入场价={self.entry_price}, 当前价={current_price}, 盈利比例={profit_pct:.2%}")
            return True
            
        return False
    
    def get_stop_loss_signal(self, current_price):
        """获取止损信号"""
        if self.check_stop_loss(current_price):
            return -1  # 止损信号为卖出
        return 0
    
    def get_take_profit_signal(self, current_price):
        """获取止盈信号"""
        if self.check_take_profit(current_price):
            # 返回部分卖出信号（负整数表示卖出数量）
            return int(-self.position * self.partial_take_profit_pct)
        return 0

class BacktestEngine:
    """回测引擎"""
    
    def __init__(self, initial_capital=100000, commission_rate=0.0003, slippage=0.001):
        """
        初始化回测引擎
        
        参数:
            initial_capital: 初始资金
            commission_rate: 交易手续费率
            slippage: 滑点比例
        """
        self.initial_capital = initial_capital
        self.commission_rate = commission_rate
        self.slippage = slippage
        self.positions = {}  # 持仓情况 {股票代码: 数量}
        self.cash = initial_capital  # 可用现金
        self.portfolio_value = initial_capital  # 总资产
        self.trades = []  # 交易记录
        self.equity_curve = []  # 净值曲线
        
    def run_backtest(self, data, strategy, stop_loss_take_profit=None):
        """
        运行回测
        
        参数:
            data: 包含价格数据的DataFrame
            strategy: 策略对象
            stop_loss_take_profit: 止损止盈对象
        """
        # 生成交易信号
        signals = strategy.generate_signals(data)
        
        # 合并信号和价格数据
        backtest_data = pd.concat([data, signals], axis=1)
        
        # 初始化持仓和现金
        self.positions = {}
        self.cash = self.initial_capital
        self.portfolio_value = self.initial_capital
        self.trades = []
        self.equity_curve = []
        
        # 遍历每一天的数据进行回测
        for i in range(1, len(backtest_data)):
            current_date = backtest_data.index[i]
            prev_date = backtest_data.index[i-1]
            
            # 获取当前价格和前一天的信号
            current_price = backtest_data.loc[current_date, 'close']
            signal = backtest_data.loc[prev_date, 'signal']
            
            # 更新止损止盈机制中的持仓信息
            if stop_loss_take_profit:
                position = self.positions.get('stock', 0)
                stop_loss_take_profit.update_position(position, current_price if position > 0 else 0)
                
                # 检查止损信号
                sl_signal = stop_loss_take_profit.get_stop_loss_signal(current_price)
                if sl_signal != 0:
                    signal = sl_signal  # 止损信号优先级最高
                    
                # 检查止盈信号（如果没有触发止损）
                if signal == 0:
                    tp_signal = stop_loss_take_profit.get_take_profit_signal(current_price)
                    if tp_signal != 0:
                        signal = tp_signal  # 部分止盈信号
            
            # 执行交易
            if signal != 0:
                self.execute_trade(
                    date=current_date,
                    price=current_price,
                    signal=signal
                )
            
            # 计算当日资产价值
            position_value = self.positions.get('stock', 0) * current_price
            self.portfolio_value = self.cash + position_value
            
            # 记录净值曲线
            self.equity_curve.append({
                'date': current_date,
                'portfolio_value': self.portfolio_value,
                'cash': self.cash,
                'position': self.positions.get('stock', 0),
                'position_value': position_value
            })
            
            # 日志记录
            if i % 100 == 0:
                logger.info(f"回测进度: {i}/{len(backtest_data)}, 日期: {current_date}, 资产: {self.portfolio_value:.2f}")
        
        # 转换净值曲线为DataFrame
        self.equity_curve = pd.DataFrame(self.equity_curve).set_index('date')
        
        return self.calculate_performance()
    
    def execute_trade(self, date, price, signal):
        """
        执行交易
        
        参数:
            date: 交易日期
            price: 交易价格
            signal: 交易信号 (1=买入, -1=卖出, 负整数=部分卖出)
        """
        current_position = self.positions.get('stock', 0)
        
        if signal > 0:  # 买入
            # 计算可买入数量（考虑手续费和滑点）
            effective_price = price * (1 + self.slippage)
            max_amount = self.cash / (1 + self.commission_rate)
            shares_to_buy = int(max_amount / effective_price)
            
            if shares_to_buy > 0:
                # 计算交易成本
                transaction_cost = shares_to_buy * effective_price * (1 + self.commission_rate)
                
                # 更新持仓和现金
                self.positions['stock'] = current_position + shares_to_buy
                self.cash -= transaction_cost
                
                # 记录交易
                self.trades.append({
                    'date': date,
                    'type': '买入',
                    'price': effective_price,
                    'shares': shares_to_buy,
                    'amount': transaction_cost,
                    'position_after': self.positions['stock']
                })
                
                logger.info(f"{date} 买入: {shares_to_buy}股, 价格: {effective_price:.2f}, 总额: {transaction_cost:.2f}")
                
        elif signal < 0:  # 卖出
            if current_position <= 0:
                return
                
            if signal == -1:  # 全部卖出
                shares_to_sell = current_position
            else:  # 部分卖出
                shares_to_sell = min(current_position, abs(signal))
                
            if shares_to_sell > 0:
                # 计算交易收入（考虑手续费和滑点）
                effective_price = price * (1 - self.slippage)
                transaction_income = shares_to_sell * effective_price * (1 - self.commission_rate)
                
                # 更新持仓和现金
                self.positions['stock'] = current_position - shares_to_sell
                self.cash += transaction_income
                
                # 记录交易
                self.trades.append({
                    'date': date,
                    'type': '卖出',
                    'price': effective_price,
                    'shares': shares_to_sell,
                    'amount': transaction_income,
                    'position_after': self.positions['stock']
                })
                
                logger.info(f"{date} 卖出: {shares_to_sell}股, 价格: {effective_price:.2f}, 总额: {transaction_income:.2f}")
    
    def calculate_performance(self):
        """计算回测绩效指标"""
        if not self.equity_curve.empty:
            # 计算每日收益率
            self.equity_curve['daily_return'] = self.equity_curve['portfolio_value'].pct_change()
            
            # 计算累计收益率
            total_return = (self.portfolio_value / self.initial_capital) - 1
            
            # 计算年化收益率
            days = (self.equity_curve.index[-1] - self.equity_curve.index[0]).days
            annual_return = (1 + total_return) ** (365 / days) - 1
            
            # 计算夏普比率
            risk_free_rate = 0.03  # 假设无风险年化收益率为3%
            daily_risk_free = (1 + risk_free_rate) ** (1/252) - 1
            excess_return = self.equity_curve['daily_return'] - daily_risk_free
            sharpe_ratio = np.sqrt(252) * excess_return.mean() / excess_return.std()
            
            # 计算最大回撤
            self.equity_curve['peak'] = self.equity_curve['portfolio_value'].cummax()
            self.equity_curve['drawdown'] = (self.equity_curve['portfolio_value'] - self.equity_curve['peak']) / self.equity_curve['peak']
            max_drawdown = self.equity_curve['drawdown'].min()
            
            # 计算胜率
            trades_df = pd.DataFrame(self.trades)
            if not trades_df.empty and 'type' in trades_df.columns and 'amount' in trades_df.columns:
                sell_trades = trades_df[trades_df['type'] == '卖出']
                if not sell_trades.empty:
                    win_trades = sell_trades[sell_trades['amount'] > 0]
                    win_rate = len(win_trades) / len(sell_trades)
                else:
                    win_rate = 0
            else:
                win_rate = 0
            
            # 计算交易次数
            trade_count = len(self.trades)
            
            # 返回绩效指标
            performance = {
                '总收益率': total_return,
                '年化收益率': annual_return,
                '夏普比率': sharpe_ratio,
                '最大回撤': max_drawdown,
                '胜率': win_rate,
                '交易次数': trade_count,
                '最终资产': self.portfolio_value,
                '回测天数': days
            }
            
            return performance
        else:
            return {}
    
    def plot_equity_curve(self):
        """绘制净值曲线"""
        if not self.equity_curve.empty:
            plt.figure(figsize=(14, 7))
            plt.plot(self.equity_curve.index, self.equity_curve['portfolio_value'])
            plt.title('回测净值曲线')
            plt.xlabel('日期')
            plt.ylabel('资产价值')
            plt.grid(True, alpha=0.3)
            
            # 设置x轴日期格式
            date_format = DateFormatter('%Y-%m-%d')
            plt.gca().xaxis.set_major_formatter(date_format)
            plt.gcf().autofmt_xdate()
            
            plt.tight_layout()
            return plt.gcf()
        else:
            print("没有净值曲线数据可供绘制")
            return None
    
    def plot_drawdown(self):
        """绘制回撤曲线"""
        if not self.equity_curve.empty and 'drawdown' in self.equity_curve.columns:
            plt.figure(figsize=(14, 7))
            plt.fill_between(self.equity_curve.index, self.equity_curve['drawdown'], 0, color='red', alpha=0.3)
            plt.title('回测回撤曲线')
            plt.xlabel('日期')
            plt.ylabel('回撤比例')
            plt.grid(True, alpha=0.3)
            
            # 设置y轴为百分比格式
            plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(1.0))
            
            # 设置x轴日期格式
            date_format = DateFormatter('%Y-%m-%d')
            plt.gca().xaxis.set_major_formatter(date_format)
            plt.gcf().autofmt_xdate()
            
            plt.tight_layout()
            return plt.gcf()
        else:
            print("没有回撤数据可供绘制")
            return None

# 示例：如何使用上述类进行回测
def run_example_backtest():
    """运行示例回测"""
    dates = pd.date_range(start='2020-01-01', end='2023-12-31', freq='B')
    np.random.seed(42)
    
    # 修复：将价格数据转换为 Pandas Series
    prices = pd.Series(
        100 * np.cumprod(1 + np.random.normal(0, 0.01, len(dates))),
        index=dates
    )
    
    data = pd.DataFrame({
        'close': prices,
        'MA5': prices.rolling(window=5).mean(),  # 现在可以正确调用 rolling 方法
        'MA20': prices.rolling(window=20).mean(),
        'DIF': prices.ewm(span=12).mean() - prices.ewm(span=26).mean(),
        'DEA': (prices.ewm(span=12).mean() - prices.ewm(span=26).mean()).ewm(span=9).mean(),
        'K': np.random.uniform(20, 80, len(dates)),
        'D': np.random.uniform(20, 80, len(dates)),
        'profit_grow_rate': np.random.normal(0.1, 0.2, len(dates)),
        'debt_to_assets': np.random.normal(0.4, 0.1, len(dates))
    }, index=dates)
    

    
    # 创建策略
    tech_strategy = TechnicalStrategy()
    fundamental_strategy = FundamentalStrategy()
    
    # 创建止损止盈机制
    sl_tp = StopLossTakeProfit(stop_loss_pct=0.05, take_profit_pct=0.15)
    
    # 创建回测引擎
    backtest_engine = BacktestEngine(initial_capital=100000, commission_rate=0.0003, slippage=0.001)
    
    # 运行回测（量价策略）
    print("运行量价策略回测...")
    tech_performance = backtest_engine.run_backtest(data, tech_strategy, sl_tp)
    tech_equity_curve = backtest_engine.equity_curve.copy()
    
    # 运行回测（财务策略）
    print("运行财务策略回测...")
    fundamental_performance = backtest_engine.run_backtest(data, fundamental_strategy, sl_tp)
    fundamental_equity_curve = backtest_engine.equity_curve.copy()
    
    # 打印绩效指标
    print("\n量价策略绩效:")
    for key, value in tech_performance.items():
        if isinstance(value, float):
            print(f"{key}: {value:.2%}" if key != '交易次数' and key != '回测天数' else f"{key}: {value}")
        else:
            print(f"{key}: {value}")
    
    print("\n财务策略绩效:")
    for key, value in fundamental_performance.items():
        if isinstance(value, float):
            print(f"{key}: {value:.2%}" if key != '交易次数' and key != '回测天数' else f"{key}: {value}")
        else:
            print(f"{key}: {value}")
    
    # 绘制对比图
    plt.figure(figsize=(14, 8))
    plt.plot(tech_equity_curve.index, tech_equity_curve['portfolio_value'] / 100000, label='量价策略')
    plt.plot(fundamental_equity_curve.index, fundamental_equity_curve['portfolio_value'] / 100000, label='财务策略')
    plt.axhline(y=1, color='black', linestyle='-', alpha=0.3)
    plt.title('策略净值曲线对比')
    plt.xlabel('日期')
    plt.ylabel('净值')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    run_example_backtest()


# In[ ]:




